package cn.iocoder.yudao.framework.crypto.core;

import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.iocoder.yudao.framework.common.exception.KnownServiceException;
import cn.iocoder.yudao.framework.common.pojo.CommonResult;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.framework.crypto.annotation.ApiResponseEncrypt;
import cn.iocoder.yudao.framework.crypto.pojo.CryptCommonResult;
import cn.iocoder.yudao.framework.crypto.pojo.CryptoProperties;
import jakarta.annotation.Nonnull;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

/**
 * <pre>
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 *  响应数据加密 Advice
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 * </pre>
 *
 * @author Zhougang
 * @author 山野羡民（1032694760@qq.com）
 * @since 2025/03/24
 * @see cn.iocoder.yudao.framework.web.core.handler.GlobalResponseBodyHandler
 */
@Slf4j
@RestControllerAdvice
@RequiredArgsConstructor
public class ApiEncryptResponseBodyAdvice implements ResponseBodyAdvice<CommonResult<?>> {
    private final CryptoProperties cryptoProperties;

    @Override
    public boolean supports(@Nonnull MethodParameter returnType, @Nonnull Class<? extends HttpMessageConverter<?>> converterType) {
        if (returnType.getMethod() == null) {
            return false;
        }
        boolean support = false;
        if (cryptoProperties.isEnable()) {
            // 功能启用
            if (returnType.hasMethodAnnotation(ApiResponseEncrypt.class)) {
                // 方法加了 @ApiResponseEncrypt 注解
                support = true;
            } else if (returnType.getContainingClass().isAnnotationPresent(ApiResponseEncrypt.class)) {
                // 类或子类加了 @ApiResponseEncrypt 注解
                support = true;
            }
        }
        if (support && cryptoProperties.isShowLog()) {
            log.info("[supports][响应结果加密 Advice 拦截处理] returnType={}", returnType);
        }
        return support;
    }

    @Override
    public CommonResult<?> beforeBodyWrite(CommonResult<?> body, @Nonnull MethodParameter returnType, @Nonnull MediaType selectedContentType, @Nonnull Class<? extends HttpMessageConverter<?>> selectedConverterType, @Nonnull ServerHttpRequest request, @Nonnull ServerHttpResponse response) {
        if (body == null) {
            return null;
        }
        // 获取客户端 AES 加密密钥
        String encryptedAesKey = request.getHeaders().getFirst(cryptoProperties.getAesHeader());
        if (StrUtil.isBlank(encryptedAesKey)) {
            log.error("[beforeBodyWrite]响应结果加密异常，缺少请求头值：{}", cryptoProperties.getAesHeader());
            throw new KnownServiceException("请求头 %s 的值为空".formatted(cryptoProperties.getAesHeader()));
        }
        try {
            // 通过 RSA 私钥解密 AES 密钥
            byte[] decryptedAesKey = SecureUtil.rsa(cryptoProperties.getPrivateKey(), cryptoProperties.getPublicKey()).decrypt(encryptedAesKey, KeyType.PrivateKey);
            // 使用 AES 密钥加密数据
            if (body.getData() == null) {
                log.error("[beforeBodyWrite]响应结果数据为空，无需加密，body：{}", body);
                return body;
            }
            byte[] plainData = JsonUtils.toJsonByte(body.getData());
            String encryptedData = SecureUtil.aes(decryptedAesKey).encryptBase64(plainData);
            if (cryptoProperties.isShowLog()) {
                log.info("[beforeBodyWrite]响应结果明文字符串为：{}，加密后：{}", body, encryptedData);
            }
            return new CryptCommonResult()
                    .setCrypto(true)
                    .setCode(body.getCode())
                    .setMsg(body.getMsg())
                    .setData(encryptedData);
        } catch (Exception e) {
            log.error("[beforeBodyWrite]响应结果加密异常，body：{}", body, e);
            throw new KnownServiceException("响应结果加密异常");
        }
    }

}
